import torch
from nets.unet import UNet
from torchvision.transforms import functional as F
from PIL import Image
import numpy as np

# 定义类别标签和对应的颜色
class_labels = ['Background', 'Bullding-Tiooded', 'Bullding-non-flodded', 'Raad-flaaded',
                'Road-non-tlooded', 'Water', 'Tree', 'Vehicle', 'Pool', 'Grass']
class_colors = [(0, 0, 0), (255, 0, 0), (165, 42, 42), (184, 134, 11),
                (128, 128, 128), (135, 206, 250), (0, 0, 139), (255, 192, 203),
                (255, 0, 0), (0, 128, 0)]

# 加载预训练模型
model_path = "path_to_model/unet_model.pth"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = UNet(num_classes=len(class_labels)).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

# 加载待预测的图像
image_path = "path_to_image/image.jpg"
input_image = Image.open(image_path).convert("RGB")

# 进行图像预处理
input_tensor = F.to_tensor(input_image).unsqueeze(0).to(device)

# 进行图像预测
with torch.no_grad():
    output = model(input_tensor)

# 处理预测结果
predicted_class_indices = torch.argmax(output, dim=1).squeeze().cpu().numpy()
predicted_colors = np.zeros((predicted_class_indices.shape[0], predicted_class_indices.shape[1], 3), dtype=np.uint8)
for i in range(len(class_labels)):
    predicted_colors[predicted_class_indices == i] = class_colors[i]

# 可选：将预测结果可视化
output_image = Image.fromarray(predicted_colors)
output_image.show()
